package com.cqie.controller;

import com.alibaba.fastjson.JSON;
import com.cqie.constant.AskAiConstant;
import com.cqie.exception.SpecialException;
import com.cqie.util.ResultTemplate;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.ObjectUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 操作人: 不吃海鲜
 * 时间: 2024/12/30 23:17
 */
@Tag(name = "AI Controller", description = "描述：Ai 接口")
@Slf4j
@RestController
@RequestMapping({"/api/v1.0/AskAi", "/api/v1/AskAi"})
public class AskAiController {

    // 请自定义自己的业务id
    private static final String requestIdTemplate = "mycompany-%d";

    /**
     * sse调用
     */
    @Operation(summary = "sse调用")
    @GetMapping("/sseInvoke")
    public ResultTemplate<String> testSseInvoke(String content) throws SpecialException {
        // 检查content是否为空
        if (ObjectUtils.isEmpty(content)) {
            throw new SpecialException("参数不能为空");
        }
        // 创建一个ChatMessage类型的列表
        List<ChatMessage> messages = new ArrayList<>();
        // 创建一个ChatMessage对象，角色为USER，内容为传入的content
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), content);
        // 将ChatMessage对象添加到列表中
        messages.add(chatMessage);
        // 生成一个请求ID，格式为当前时间的毫秒数
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());

        // 创建一个ChatTool类型的列表
        List<ChatTool> chatToolList = new ArrayList<>();
        ChatTool chatTool = new ChatTool();
        // 设置ChatTool的类型为FUNCTION
        chatTool.setType(ChatToolType.FUNCTION.value());
        ChatFunctionParameters chatFunctionParameters = new ChatFunctionParameters();
        // 设置ChatFunctionParameters的类型为"object"
        chatFunctionParameters.setType("object");
        // 创建一个Map对象，用于存储属性
        Map<String, Object> properties = new HashMap<>();
        properties.put("departure", new HashMap<String, Object>() {{
            // 设置departure属性的类型为string
            put("type", "string");
            // 设置departure属性的描述
            put("description", "出发城市或车站");
        }});
        properties.put("destination", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "目的地城市或车站");
        }});
        properties.put("date", new HashMap<String, Object>() {{
            put("type", "string");
            put("description", "要查询的车次日期");
        }});


        List<String> required = new ArrayList<>();
        required.add("departure");
        required.add("destination");
        required.add("date");

        // 设置ChatFunctionParameters的属性
        chatFunctionParameters.setProperties(properties);
        ChatFunction chatFunction = ChatFunction.builder()
                .name("query_medical_info")
                .description("根据用户提供的信息，查询对应的车次")
                .parameters(chatFunctionParameters)
                .required(required)
                .build();

        // 设置ChatTool的函数
        chatTool.setFunction(chatFunction);
        chatToolList.add(chatTool);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4) // 设置模型为Constants.ModelChatGLM4
                .stream(Boolean.TRUE) // 设置流为true
                .messages(messages)// 设置消息列表
                .requestId(requestId) // 设置请求ID
                .tools(chatToolList) // 设置工具列表
                .toolChoice("auto") // 设置工具选择为"auto"
                .build();

        // 调用模型API，获取响应
        ModelApiResponse sseModelApiResp = AskAiConstant.client.invokeModelApi(chatCompletionRequest);
        StringBuilder responseBuilder = new StringBuilder();
        // 检查响应是否成功
        if (sseModelApiResp.isSuccess()) {
            // 用于标记是否是第一次处理
            AtomicBoolean isFirst = new AtomicBoolean(true);
            // 将响应流转换为ChatMessageAccumulator
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    // 对每个accumulator进行处理
                    .doOnNext(accumulator -> {
                        // 如果是第一次处理，则不添加任何内容
                        if (isFirst.getAndSet(false)) {
                            responseBuilder.append("");
                        }

                        // 如果accumulator的delta不为空且包含工具调用
                        if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                            String jsonString =
                                    AskAiConstant.mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                            responseBuilder.append("工具调用: ").append(jsonString);
                        }

                        // 如果accumulator的delta不为空且包含内容
                        if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                            responseBuilder.append(accumulator.getDelta().getContent());
                        }
                    })
                    // 在完成时添加"完成"到响应内容中
                    .doOnComplete(() -> responseBuilder.append("\n完成"))
                    // 获取最后一个元素
                    .lastElement()
                    // 阻塞获取结果
                    .blockingGet();

            // 将收集到的响应内容设置到ModelData中
            ModelData modelData = new ModelData();
            modelData.setModel(responseBuilder.toString());
            sseModelApiResp.setData(modelData);
        }
        System.out.println("模型输出:" + JSON.toJSONString(sseModelApiResp));
        return ResultTemplate.Success(sseModelApiResp.getData());
    }

    public static Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable.map(chunk -> {
            return new ChatMessageAccumulator(chunk.getChoices().get(0).getDelta(), null, chunk.getChoices().get(0),
                    chunk.getUsage(), chunk.getCreated(), chunk.getId());
        });
    }
}
